PROBER = "TCPIP0::192.168.0.2::800::SOCKET"
SMU = "GPIB0::17::INSTR"
SWITCH = "GPIB0::22::INSTR"

import pyvisa
rm = pyvisa.ResourceManager()

linkCheck = True
try:
    prober = rm.open_resource(PROBER)
    prober.read_termination = '\r\n'
    prober.query("*IDN?")
except Exception as e:
    print("MPI Prober Station Link Error!")
    prober = None
    linkCheck =  False
try:
    B1500a = rm.open_resource(SMU)
    B1500a.query("*IDN?")
except Exception as e:
    print(f"B1500 SMU Link Error!")
    linkCheck =  False
    B1500a = None
try:
    B2201 = rm.open_resource(SWITCH)
    B2201.query("*IDN?")
except Exception as e:
    print("B2201 Switching Matrix Link Error!")
    linkCheck =  False
    B2201 = None

def contact(prober=prober):
    prober.write("move_chuck_contact")
def sep(prober=prober):
    prober.write("move_chuck_separation")
def nextdie(prober=prober):
    prober.write("map:step_next_die")
def subdie(prober,x,y,xoffset=0,yoffset=0):
    prober.write(f'map:step_specify_die {x},{y},{xoffset},{yoffset}')

def matrix(stat = 4,b2201=B2201):
    import B2200
    import time
    #1,1234-1234
    #2,1234-5678
    #3,1234-9-12
    #4,openall
    B2200.Load_Path(b2201,stat)
    
def matrixx(pos):
    SWITCH = "GPIB0::22::INSTR"

    import pyvisa
    rm = pyvisa.ResourceManager()
    B2201 = rm.open_resource(SWITCH)

    #switching matrix
    import B2200
    import time
    #1,1234-1234
    #2,1234-5678
    #3,1234-9-12
    #4,openall
    B2200.Load_Path(B2201,pos)
    #time.sleep(2)
    #B2200.Load_Path(B2201,2)
    #time.sleep(2)
    #B2200.Load_Path(B2201,3)
    #time.sleep(2)
    #B2200.Load_Path(B2201,4)

    rm.close()
matrixx(1)


#b1500a smu
def transCurve(DeviceType = 'P',
               useDefault=True,
               startV=-2, endV=0.5,Vds=-2,
               chnSource='1',chnDrain='2',chnGate='3',
               B1500a=B1500a,
               ADC='AIT 1,2,1',
               direction='Single Linear',
               reverse=False,
               points=101,
               t_hold=1.0,s_delay=0.0,t_delay=0.0,**kwarg):

    if useDefault:
        if DeviceType == 'P':
            startV, endV, Vds, chnSource = -2, 0.5, -2, '1'
            chnDrain='2'
            chnGate='3'
        elif DeviceType == 'N':
            startV, endV, Vds, chnSource = 2, -0.5, 2, '4'
            chnDrain='2'
            chnGate='3'
    if reverse:
        startV,endV = endV,startV

    import pyvisa
    rm = pyvisa.ResourceManager()
    SMU = "GPIB0::17::INSTR"
    B1500a = rm.open_resource(SMU)
    B1500a.write('*RST')
    B1500a.write('CN')

    #--------------------IV_Sweep函数测试参数示例（双变量扫描）------------------#
    Unit1=chnGate                    #扫描变量1对应的物理端口为4
    VName1='VGate'                   #扫描变量1的名字为VD，注意电压模式的端口命名为V....,电流模式端口命名为I...
    Direction1=direction    #扫描变量1的方向为Double，并且加电压模式为Linear（更多模式见34）
    Start1=startV                   #扫描变量1的起始电压为-3 V
    end1=endV                     #扫描变量1的终止电压为3 V
    nop1=points                      #扫描变量1的点数目为201（Double的话，实际数据点为201*2）
    Comp1=1e-2                  #扫描变量1的限流值
    range1='Auto'         #扫描变量1的测量Range

    VAR_setup=[Unit1,VName1,Direction1,Start1,end1,nop1,Comp1,range1] #打包扫描变量，如果单变量扫描，不用第二行

    Measurement_setup = []

    Unit=chnDrain                  #测量变量1的物理端口为
    VName='VDrain'                  #测量变量1的名称，注意测电流是固定电压名称为V...，测电压是固定电流，名称是I...
    VorI=Vds                 #测量变量1的固定电压为0 V
    Comp=0.1                 #测量变量1的限流值
    range='Auto'                #测量变量1的Range
    Measurement_setup  += [Unit,VName,VorI,Comp,range] #打包测量变量，有几个打包几个

    Unit=chnSource                   #测量变量1的物理端口为
    VName='VSource'                  #测量变量1的名称，注意测电流是固定电压名称为V...，测电压是固定电流，名称是I...
    VorI=0                 #测量变量1的固定电压为0 V
    Comp=0.1                  #测量变量1的限流值
    range='Auto'                #测量变量1的Range
    Measurement_setup  += [Unit,VName,VorI,Comp,range] #打包测量变量，有几个打包几个


    #print(Measurement_setup)
    B1500a.timeout = 60000
    B1500a.write('*RST')
    B1500a.write('CN 1,2,3,4')
    import B1500
    res = B1500.IV_sweept(B1500a,VAR_setup,Measurement_setup,ADC=ADC,t_hold=t_hold,s_delay=s_delay,t_delay=t_delay)#AIT 1,2,1 20ms 1,1,125 10ms

    import pandas
    df = pandas.DataFrame(res[1:],columns=['GateV','GateI','DrainV','DrainI','SourceV','SourceI'],dtype=float)
    B1500a.close()
    rm.close()
    return df

PMOSTRANS = {"name":'PMOS',"startV":-2, "endV":0.5,"Vds":-2,
            "chnSource":'1',"chnDrain":'2',"chnGate":'3'}
NMOSTRANS = {"name":'NMOS',"startV":2, "endV":-0.5,"Vds":2,
            "chnSource":'4',"chnDrain":'2',"chnGate":'3'}
FZ1 = {"dies":58,"lines":3,"lineOffset":0.342}
FZ2 = {"dies":50,"lines":6,"lineOffset":0.249}

from functools import partial
transCurveS = partial(transCurve,'P',False)

def autoTest(tests=[PMOSTRANS,NMOSTRANS],
             dies=58,
             lines=3,
             lineOffset=0.342,
             positions=[1,2,3]):
    if not linkCheck:
        return 
    import pandas as pd 
    from datetime import datetime
    tStamp = datetime.strftime(datetime.now(),"%Y%m%d_%H%M%S")
    import os
    os.mkdir(f'autoData/{tStamp}')
    xlsFiles = {}
    from itertools import product
    for line,position,test in product(range(1,lines+1),positions,tests):
        fn = f'{test['name']}-LINE{line}-{str(position)}'
        xlsFiles[fn] = pd.ExcelWriter(f'autoData/{tStamp}/{fn}.xlsx')

    import matplotlib.pyplot as plt
    import numpy as np
    import time

    plt.style.use('default')
    plt.ion() #开启interactive mode 成功的关键函数
    plt.cla()
    plt.xlabel('GateV (V)')
    plt.ylabel('SourceI (A)')
    plt.yscale('log')
    plt.title('StarAutoProbe')

    import time
    for i in range(dies):
        nextdie()
        for line in lines:
            if line!=1:
                prober.write(f'move_chuck_xy Relative,0,{lineOffset}')
            time.sleep(0.3)
            contact()
            time.sleep(0.3)
            for position in positions:
                matrix(position)
                time.sleep(0.2)
                for test in tests:
                    fn = f'{test['name']}-LINE{line}-{str(position)}'
                    print('\r'+str(i+1)+f'/{dies}:{fn}',end='')
                    try:
                        df = transCurveS(**test)
                        plt.title("StarAutoProbe: "+str(i+1)+f'/{dies}:{fn}')
                        plt.plot(df['GateV'],np.abs(df['SourceI']))
                        plt.pause(0.1)
                        df.to_excel(xlsFiles[fn],index=False,sheet_name=f"Run{i}")
                    except:
                        print("\r"+str(i+1)+f'/{dies}:{fn}'+"error, skipped!")
                        try:
                            for line,position,test in product(range(1,lines+1),positions,tests):
                                fn = f'{test['name']}-LINE{line}-{str(position)}'
                                xlsFiles[fn].close()
                            print("Data saved.")
                        except:
                            pass
                matrix(4)
            print('\r'+str(i+1)+'/50',end='')
            if not (i == {dies}-1 and line == lines-1):
                sep()
    plt.title('StarAutoProbe: Complete!')
    plt.pause(0.1)

    for line,position,test in product(range(1,lines+1),positions,tests):
        fn = f'{test['name']}-LINE{line}-{str(position)}'
        xlsFiles[fn].close()


def IVCurve(startV=-2, endV=2,
               chn1='1',chn2='2',
               B1500a=B1500a,
               ADC='AIT 1,2,1',
               direction='Single Linear',
               points=201,
               t_hold=1.0,s_delay=0.0,t_delay=0.0):

    import pyvisa
    rm = pyvisa.ResourceManager()
    SMU = "GPIB0::17::INSTR"
    B1500a = rm.open_resource(SMU)
    B1500a.write('*RST')
    B1500a.write('CN')

    #--------------------IV_Sweep函数测试参数示例（双变量扫描）------------------#
    Unit1=chn1                   #扫描变量1对应的物理端口为4
    VName1='VCH1'                   #扫描变量1的名字为VD，注意电压模式的端口命名为V....,电流模式端口命名为I...
    Direction1=direction    #扫描变量1的方向为Double，并且加电压模式为Linear（更多模式见34）
    Start1=startV                   #扫描变量1的起始电压为-3 V
    end1=endV                     #扫描变量1的终止电压为3 V
    nop1=points                      #扫描变量1的点数目为201（Double的话，实际数据点为201*2）
    Comp1=1e-2                  #扫描变量1的限流值
    range1='Auto'         #扫描变量1的测量Range

    VAR_setup=[Unit1,VName1,Direction1,Start1,end1,nop1,Comp1,range1] #打包扫描变量，如果单变量扫描，不用第二行

    Measurement_setup = []

    Unit=chn2                 #测量变量1的物理端口为
    VName='VCH2'                  #测量变量1的名称，注意测电流是固定电压名称为V...，测电压是固定电流，名称是I...
    VorI=0                 #测量变量1的固定电压为0 V
    Comp=0.1                 #测量变量1的限流值
    range='Auto'                #测量变量1的Range
    Measurement_setup  += [Unit,VName,VorI,Comp,range] #打包测量变量，有几个打包几个

    #print(Measurement_setup)
    B1500a.timeout = 60000
    B1500a.write('*RST')
    B1500a.write('CN 1,2,3,4')
    import B1500
    res = B1500.IV_sweept(B1500a,VAR_setup,Measurement_setup,ADC=ADC,t_hold=t_hold,s_delay=s_delay,t_delay=t_delay)#AIT 1,2,1 20ms 1,1,125 10ms

    import pandas
    df = pandas.DataFrame(res[1:],columns=['Voltage','Current','GNDV','GNDI'],dtype=float)
    B1500a.close()
    rm.close()
    return df

#b1500a smu
def inverter(B1500a=B1500a,
        ADC='AIT 1,2,1',
        direction='Double Linear',
        points=101,
        t_hold=0.0,s_delay=0.0,t_delay=0.0):
    B1500a.write('*RST')
    B1500a.write('CN')

    #--------------------IV_Sweep函数测试参数示例（双变量扫描）------------------#
    Unit1='3'                     #扫描变量1对应的物理端口为4
    VName1='VVin'                   #扫描变量1的名字为VD，注意电压模式的端口命名为V....,电流模式端口命名为I...
    Direction1=direction    #扫描变量1的方向为Double，并且加电压模式为Linear（更多模式见34）
    Start1=0                   #扫描变量1的起始电压为-3 V
    end1=2.0                      #扫描变量1的终止电压为3 V
    nop1=points                    #扫描变量1的点数目为201（Double的话，实际数据点为201*2）
    Comp1=1e-2                  #扫描变量1的限流值
    range1='Auto'         #扫描变量1的测量Range

    VAR_setup=[Unit1,VName1,Direction1,Start1,end1,nop1,Comp1,range1] #打包扫描变量，如果单变量扫描，不用第二行

    Measurement_setup = []
    Unit='1'                    #测量变量1的物理端口为
    VName='VVdd'                  #测量变量1的名称，注意测电流是固定电压名称为V...，测电压是固定电流，名称是I...
    VorI=2.0                 #测量变量1的固定电压为0 V
    Comp=1e-2                  #测量变量1的限流值
    range='Auto'                #测量变量1的Range
    Measurement_setup  += [Unit,VName,VorI,Comp,range] #打包测量变量，有几个打包几个

    Unit='4'                    #测量变量1的物理端口为
    VName='VGND'                  #测量变量1的名称，注意测电流是固定电压名称为V...，测电压是固定电流，名称是I...
    VorI=0.0                 #测量变量1的固定电压为0 V
    Comp=1e-2                  #测量变量1的限流值
    range='Auto'                #测量变量1的Range
    Measurement_setup  += [Unit,VName,VorI,Comp,range] #打包测量变量，有几个打包几个

    Unit='2'                    #测量变量1的物理端口为
    VName='IVout'                  #测量变量1的名称，注意测电流是固定电压名称为V...，测电压是固定电流，名称是I...
    VorI=0.0                 #测量变量1的固定电压为0 V
    Comp=5                  #测量变量1的限流值
    range='Auto'                #测量变量1的Range
    Measurement_setup  += [Unit,VName,VorI,Comp,range] #打包测量变量，有几个打包几个
    #print(Measurement_setup)
    B1500a.timeout = 60000
    B1500a.write('*RST')
    B1500a.write('CN 1,2,3,4')
    import B1500
    res = B1500.IV_sweept(B1500a,VAR_setup,Measurement_setup,ADC=ADC,t_hold=t_hold,s_delay=s_delay,t_delay=t_delay)

    import pandas
    df = pandas.DataFrame(res[1:],columns=['Vin','Iin','Vdd','Idd','VGND','IGND',"Iout",'Vout'],dtype=float)
    B1500a.close()
    return df

def closeIns():
    if prober:
        prober.close()
    if B1500a:
        B1500a.close()
    if B2201:
        B2201.close()
    rm.close()    